#include "WebSocketSessionHandler.h"
#include "Poco/Net/NetException.h"

WebSocketSessionHandler::WebSocketSessionHandler(const std::string& targetHost, int targetPort, const std::string& targetPath)
    : _targetHost(targetHost), _targetPort(targetPort), _targetPath(targetPath) {}

void WebSocketSessionHandler::handleRequest(Poco::Net::HTTPServerRequest& request, Poco::Net::HTTPServerResponse& response) {
    try {
        std::cout << "Proxying WebSocket Request: " << request.getURI() << " to ws://" << _targetHost << ":" << _targetPort << _targetPath << std::endl;

        // 1. Establish WebSocket connection with the client
        Poco::Net::WebSocket clientWs(request, response);
        clientWs.setReceiveTimeout(Poco::Timespan(0, 500000)); // 0.5 seconds

        // 2. Establish WebSocket connection with the backend server
        Poco::Net::HTTPClientSession backendSession(_targetHost, _targetPort);
        Poco::Net::HTTPRequest backendRequest(Poco::Net::HTTPRequest::HTTP_GET, _targetPath, Poco::Net::HTTPMessage::HTTP_1_1);
        backendRequest.set("Host", _targetHost);
        // backendRequest.set("Upgrade", "websocket");
        // backendRequest.set("Connection", "Upgrade");
        // backendRequest.set("Sec-WebSocket-Version", "13");
        // backendRequest.set("Sec-WebSocket-Key", Poco::Net::WebSocket::generateKey()); // Generate a new key for backend

        // Optional: Copy client's origin header to backend if needed for security
        if (request.has("Origin")) {
            backendRequest.set("Origin", request.get("Origin"));
        }

        Poco::Net::HTTPResponse backendResponse;
        // The Poco::Net::WebSocket constructor for client-side connections
        // automatically adds the necessary WebSocket handshake headers (including Sec-WebSocket-Key, Upgrade, Connection, Sec-WebSocket-Version).
        Poco::Net::WebSocket backendWs(backendSession, backendRequest, backendResponse);
        backendWs.setReceiveTimeout(Poco::Timespan(0, 500000)); // 0.5 seconds

        // Check if backend WebSocket handshake was successful
        if (backendResponse.getStatus() != Poco::Net::HTTPResponse::HTTP_SWITCHING_PROTOCOLS) {
            std::cerr << "Backend WebSocket handshake failed: " << backendResponse.getStatus() << " " << backendResponse.getReason() << std::endl;
            response.setStatusAndReason(Poco::Net::HTTPResponse::HTTP_BAD_GATEWAY);
            std::string str = "Backend WebSocket handshake failed.";
            response.send().write(str.c_str(), str.size());
            return;
        }

        // 3. Start proxying data between client and backend
        WebSocketProxyWorker clientToBackend(clientWs, backendWs);
        WebSocketProxyWorker backendToClient(backendWs, clientWs);

        Poco::Thread thread1;
        Poco::Thread thread2;

        thread1.start(clientToBackend);
        thread2.start(backendToClient);

        thread1.join(); // Wait for clientToBackend to finish (client closes or error)
        thread2.join(); // Wait for backendToClient to finish

        std::cout << "WebSocket proxy session ended for " << request.getURI() << std::endl;

    } catch (Poco::Net::WebSocketException& exc) {
        // WebSocket handshake failed or protocol error
        std::cerr << "WebSocket Handshake/Protocol Error: " << exc.displayText() << std::endl;
        response.setStatusAndReason(Poco::Net::HTTPResponse::HTTP_BAD_REQUEST);
        std::string str = "WebSocket Error: " + exc.displayText();
        response.send().write(str.c_str(), str.size());
    } catch (Poco::Net::NetException& exc) {
        // Network related errors (connection refused, timeout etc.)
        std::cerr << "WebSocket Proxy Network Error: " << exc.displayText() << std::endl;
        response.setStatusAndReason(Poco::Net::HTTPResponse::HTTP_BAD_GATEWAY);
        std::string str = "WebSocket Proxy Network Error: " + exc.displayText();
        response.send().write(str.c_str(), str.size());
    } catch (Poco::Exception& exc) {
        std::cerr << "WebSocket Proxy Generic Error: " << exc.displayText() << std::endl;
        response.setStatusAndReason(Poco::Net::HTTPResponse::HTTP_INTERNAL_SERVER_ERROR);
        std::string str = "WebSocket Proxy Error: " + exc.displayText();
        response.send().write(str.c_str(), str.size());
    }
}

// WebSocketProxyWorker Implementation
WebSocketSessionHandler::WebSocketProxyWorker::WebSocketProxyWorker(Poco::Net::WebSocket& clientWs, Poco::Net::WebSocket& backendWs)
    : _clientWs(clientWs), _backendWs(backendWs), _stopRequested(false), _logger(Poco::Logger::get("WebSocketProxyWorker")) {}

void WebSocketSessionHandler::WebSocketProxyWorker::run() {
    char buffer[4096];
    int flags;
    int n;

    while (!_stopRequested) {
        try {
            n = _clientWs.receiveFrame(buffer, sizeof(buffer), flags);
            if (n > 0) {
                _backendWs.sendFrame(buffer, n, flags);
            } else if (flags & Poco::Net::WebSocket::FRAME_OP_CLOSE) {
                // Client initiated close, propagate to backend
                _logger.information("Client sent WebSocket close frame.");
                _backendWs.sendFrame(buffer, n, Poco::Net::WebSocket::FRAME_OP_CLOSE | flags);
                _stopRequested = true;
            }
        } catch (Poco::TimeoutException&) {
            // No data received within timeout, just continue loop
        } catch (Poco::Net::WebSocketException& exc) {
            _logger.error("WebSocket proxy error: " + exc.displayText());
            _stopRequested = true;
        } catch (Poco::Net::NetException& exc) {
            _logger.error("Network error during WebSocket proxy: " + exc.displayText());
            _stopRequested = true;
        } catch (Poco::Exception& exc) {
            _logger.error("Generic error during WebSocket proxy: " + exc.displayText());
            _stopRequested = true;
        }
    }
    _logger.information("WebSocket proxy worker stopped.");
}

void WebSocketSessionHandler::WebSocketProxyWorker::stop() {
    _stopRequested = true;
}